#!/usr/bin/env python3
# coding: utf-8
# Copyright (c) 2025 Huawei Technologies Co., Ltd.
# This file is a part of the CANN Open Software.
# Licensed under CANN Open Software License Agreement Version 1.0 (the "License").
# Please refer to the License for details. You may not use this file except in compliance with the License.
# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE.
# See LICENSE in the root of the software repository for the full text of the License.
# ======================================================================================================================
import random
import torch
from typing import Union, List

from atk.case_generator.generator.generate_types import GENERATOR_REGISTRY
from atk.case_generator.generator.base_generator import CaseGenerator
from atk.configs.case_config import InputCaseConfig, CaseConfig

@GENERATOR_REGISTRY.register("ascend_generate_grouped_matmul_swiglu_quant_v2")
class AscendGroupedMatmulSwigluQuant(CaseGenerator):
    def __init__(self, config):
        super().__init__(config)

    def after_input_config(
            self,
            index: int,
            input_case: Union[InputCaseConfig, List[InputCaseConfig]]
    ) -> Union[InputCaseConfig, List[InputCaseConfig]]:

        return input_case

    def after_case_config(self, case_config: CaseConfig) -> CaseConfig:
        x = case_config.inputs[0]
        weight = case_config.inputs[1][0]
        weightScale = case_config.inputs[2][0]
        weightAssistMatrix = case_config.inputs[3][0]
        bias = case_config.inputs[4]
        xScale = case_config.inputs[5]
        smoothScale = case_config.inputs[6]
        groupList = case_config.inputs[7]
        dequantMode = case_config.inputs[8]
        dequantDtype = case_config.inputs[9]
        quantMode = case_config.inputs[10]
        groupListType = case_config.inputs[11]
        tuningConfigOptional = case_config.inputs[12]
        print("before adjust:", x.shape, weight.shape, weightScale.shape, xScale.shape)
        M, E, K, N = x.shape[0], weight.shape[0], weight.shape[1], weight.shape[2]
        print("original shape:", M, E, K, N)
        #K轴不能大于等于65536
        if K >= 65536 or K % 16 != 0:
            K = 65536 - 16
            K = (K + 15) // 16 * 16
        #N轴不能大于10240
        if N > 10240 or N % 32 != 0:
            N = 10240
            N = (N + 31) // 32 * 32
        
        x.shape[1] = K
        weight.shape[1] = K
        weight.shape[2] = N
            
        weightScale.shape[0] = E
        weightScale.shape[1] = N
        xScale.shape[0] = M
        groupList.shape[0] = E
        print("dequantMode:",dequantMode)
        print("dequantDtype:",dequantDtype)
        print("quantMode:",quantMode)
        print("groupListType:",groupListType)
        print("tuningConfigOptional:",tuningConfigOptional)
        print("after adjust:", x.shape, weight.shape, weightScale.shape, xScale.shape)
        return case_config